#!/bin/bash

#############################################
# SSH密钥自动化部署脚本
# 功能：
# 1. 从指定URL下载公钥
# 2. 配置authorized_keys
# 3. 禁用密码登录，只允许密钥登录
# 4. 重启SSH服务
#
# 使用方法：
# 1. 交互式运行（推荐）：
#    bash -c "$(curl -fsSL https://gitee.com/xledoo/deploy-sshkey/raw/master/deploy_sshkey.sh)"
#
# 2. 非交互式运行（通过curl管道）：
#    curl -fsSL https://gitee.com/xledoo/deploy-sshkey/raw/master/deploy_sshkey.sh | sudo bash -s -- --yes
#
# 3. 非交互式运行（下载后执行）：
#    bash deploy_sshkey.sh --yes
#
# 参数说明：
#   --user=<username>     要配置SSH密钥的用户名（默认: root）
#   --yes, -y             自动确认所有提示，跳过交互
#   --help, -h            显示帮助信息
#
# 注意：
#   - 本脚本仅处理 SSH 登录密钥部署
#   - 如需配置 Git，请使用 setup_git.sh 脚本
#############################################

set -e  # 遇到错误立即退出

# 配置区域 - 请修改为您的公钥下载地址
PUBLIC_KEY_URL="https://gitee.com/xledoo/deploy-sshkey/raw/master/id_ed25519.pub"

# 颜色输出
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
NC='\033[0m' # No Color

# 参数变量
TARGET_USER="root"
AUTO_YES=false
NON_INTERACTIVE=false

# 日志函数
log_info() {
    echo -e "${GREEN}[INFO]${NC} $1"
}

log_warn() {
    echo -e "${YELLOW}[WARN]${NC} $1"
}

log_error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

# 显示帮助信息
show_help() {
    cat << EOF
SSH密钥自动化部署脚本

使用方法:
  bash deploy_sshkey.sh [选项]
  curl -fsSL https://gitee.com/xledoo/deploy-sshkey/raw/master/deploy_sshkey.sh | sudo bash -s -- [选项]

选项:
  --user=<username>     要配置SSH密钥的用户名（默认: root）
  --yes, -y             自动确认所有提示，启用非交互模式
  --help, -h            显示此帮助信息

功能:
  - 下载 SSH 公钥并配置到 authorized_keys
  - 禁用密码登录，只允许密钥登录
  - 重启 SSH 服务使配置生效

注意:
  - 本脚本仅处理 SSH 登录密钥部署
  - 如需配置 Git 和代码托管密钥，请使用 setup_git.sh 脚本

示例:
  # 交互式部署（推荐，无交互问题）
  bash -c "$(curl -fsSL https://gitee.com/xledoo/deploy-sshkey/raw/master/deploy_sshkey.sh)"

  # 非交互式部署（通过curl）
  curl -fsSL https://gitee.com/xledoo/deploy-sshkey/raw/master/deploy_sshkey.sh | sudo bash -s -- --yes

  # 非交互式部署（下载后执行）
  sudo bash deploy_sshkey.sh --yes

  # 为指定用户配置
  sudo bash deploy_sshkey.sh --user=ubuntu --yes

EOF
    exit 0
}

# 解析命令行参数
parse_arguments() {
    while [[ $# -gt 0 ]]; do
        case $1 in
            --user=*)
                TARGET_USER="${1#*=}"
                shift
                ;;
            --yes|-y)
                AUTO_YES=true
                NON_INTERACTIVE=true
                shift
                ;;
            --help|-h)
                show_help
                ;;
            *)
                log_error "未知参数: $1"
                echo "使用 --help 查看帮助信息"
                exit 1
                ;;
        esac
    done
}

# 确保脚本可以接收交互式输入（仅在交互模式下）
setup_interactive_input() {
    if [[ "$NON_INTERACTIVE" == false && ! -t 0 ]]; then
        if ( : < /dev/tty ) 2>/dev/null; then
            exec < /dev/tty
        fi
    fi
}

# 检查是否以root权限运行
check_root() {
    if [[ $EUID -ne 0 ]]; then
        log_error "此脚本必须以root权限运行"
        exit 1
    fi
    log_info "Root权限检查通过"
}

# 检查公钥URL是否已配置
check_config() {
    if [[ "$PUBLIC_KEY_URL" == "https://example.com/your-public-key.pub" ]]; then
        log_error "请先修改脚本中的PUBLIC_KEY_URL变量为实际的公钥下载地址"
        exit 1
    fi
}

# 备份SSH配置文件
backup_ssh_config() {
    local backup_file="/etc/ssh/sshd_config.backup.$(date +%Y%m%d%H%M%S)"
    cp /etc/ssh/sshd_config "$backup_file"
    log_info "SSH配置文件已备份到: $backup_file"
}

# 配置指定用户的authorized_keys
setup_authorized_keys() {
    local username=$1
    local user_home=$(eval echo ~$username)
    local ssh_dir="$user_home/.ssh"
    local auth_keys="$ssh_dir/authorized_keys"

    log_info "正在为用户 $username 配置SSH密钥..."

    # 创建.ssh目录
    if [[ ! -d "$ssh_dir" ]]; then
        mkdir -p "$ssh_dir"
        log_info "创建目录: $ssh_dir"
    fi

    # 下载公钥
    log_info "从 $PUBLIC_KEY_URL 下载公钥..."
    local temp_key="/tmp/downloaded_public_key.pub"

    if command -v curl &> /dev/null; then
        curl -fsSL "$PUBLIC_KEY_URL" -o "$temp_key"
    elif command -v wget &> /dev/null; then
        wget -q "$PUBLIC_KEY_URL" -O "$temp_key"
    else
        log_error "未找到curl或wget命令，无法下载公钥"
        exit 1
    fi

    # 验证下载的文件不为空
    if [[ ! -s "$temp_key" ]]; then
        log_error "下载的公钥文件为空"
        rm -f "$temp_key"
        exit 1
    fi

    log_info "公钥下载成功"

    # 备份现有的authorized_keys（如果存在）
    if [[ -f "$auth_keys" ]]; then
        cp "$auth_keys" "$auth_keys.backup.$(date +%Y%m%d%H%M%S)"
        log_warn "现有的authorized_keys已备份"
    fi

    # 将公钥写入authorized_keys
    cat "$temp_key" > "$auth_keys"
    rm -f "$temp_key"

    # 设置正确的权限
    chown -R $username:$username "$ssh_dir"
    chmod 700 "$ssh_dir"
    chmod 600 "$auth_keys"

    log_info "用户 $username 的SSH密钥配置完成"
    log_info "  - authorized_keys路径: $auth_keys"
    log_info "  - 目录权限: 700"
    log_info "  - 文件权限: 600"
}

# 配置SSH服务器禁用密码登录
configure_sshd() {
    log_info "正在配置SSH服务器..."

    local sshd_config="/etc/ssh/sshd_config"

    # 备份配置文件
    backup_ssh_config

    # 修改配置项的函数
    update_ssh_option() {
        local option=$1
        local value=$2

        # 如果选项已存在（无论是否注释），替换它
        if grep -q "^#*${option}" "$sshd_config"; then
            sed -i "s/^#*${option}.*/${option} ${value}/" "$sshd_config"
        else
            # 如果选项不存在，添加到文件末尾
            echo "${option} ${value}" >> "$sshd_config"
        fi
    }

    # 配置关键安全选项
    update_ssh_option "PubkeyAuthentication" "yes"
    update_ssh_option "PasswordAuthentication" "no"
    update_ssh_option "ChallengeResponseAuthentication" "no"
    update_ssh_option "UsePAM" "no"
    update_ssh_option "PermitRootLogin" "prohibit-password"
    update_ssh_option "PermitEmptyPasswords" "no"

    log_info "SSH配置已更新:"
    log_info "  - 公钥认证: 启用"
    log_info "  - 密码认证: 禁用"
    log_info "  - 质询响应认证: 禁用"
    log_info "  - PAM认证: 禁用"
    log_info "  - Root密码登录: 禁用"
}

# 测试SSH配置
test_sshd_config() {
    log_info "测试SSH配置文件..."

    if sshd -t; then
        log_info "SSH配置文件语法正确"
        return 0
    else
        log_error "SSH配置文件语法错误"
        return 1
    fi
}

# 重启SSH服务
restart_sshd() {
    log_info "正在重启SSH服务..."

    # 检测系统使用的init系统
    if command -v systemctl &> /dev/null; then
        systemctl restart sshd || systemctl restart ssh
        log_info "SSH服务已通过systemd重启"
    elif command -v service &> /dev/null; then
        service sshd restart || service ssh restart
        log_info "SSH服务已通过service重启"
    else
        log_error "无法检测到服务管理器"
        return 1
    fi
}

# 显示配置摘要和警告
show_summary() {
    echo ""
    echo "========================================="
    log_info "SSH密钥配置完成!"
    echo "========================================="
    echo ""
    log_warn "重要提示:"
    echo "  1. 请确保您的私钥已经保存在本地"
    echo "  2. 请在关闭当前SSH会话前，使用新窗口测试密钥登录"
    echo "  3. 密码登录已被禁用，请妥善保管私钥"
    echo "  4. 如果密钥登录失败，可能需要从控制台访问服务器恢复配置"
    echo ""
    echo "配置备份位置:"
    echo "  - SSH配置: /etc/ssh/sshd_config.backup.*"
    echo "  - authorized_keys: ~/.ssh/authorized_keys.backup.*"
    echo "========================================="
}

# 主函数
main() {
    # 解析命令行参数
    parse_arguments "$@"

    # 设置交互式输入
    setup_interactive_input

    log_info "开始SSH密钥自动化部署..."
    if [[ "$NON_INTERACTIVE" == true ]]; then
        log_info "运行模式: 非交互式"
    else
        log_info "运行模式: 交互式"
    fi
    echo ""

    # 检查root权限
    check_root

    # 检查配置
    check_config

    echo ""

    # 获取目标用户
    local target_user="$TARGET_USER"

    # 如果是交互模式且未提供用户，则提示输入
    if [[ "$NON_INTERACTIVE" == false ]]; then
        read -p "请输入要配置SSH密钥的用户名 [默认: $target_user]: " input_user
        target_user=${input_user:-$target_user}
        echo ""
    fi

    # 验证用户存在
    if ! id "$target_user" &>/dev/null; then
        log_error "用户 $target_user 不存在"
        exit 1
    fi

    log_warn "警告: 此操作将:"
    echo "  1. 为用户 $target_user 配置SSH登录公钥"
    echo "  2. 禁用SSH密码登录"
    echo "  3. 重启SSH服务"
    echo ""
    log_warn "提示: 如需配置 Git 和代码托管密钥，请使用 setup_git.sh 脚本"
    echo ""

    if [[ "$AUTO_YES" == false ]]; then
        read -p "确认继续? (yes/no): " confirm

        if [[ "$confirm" != "yes" ]]; then
            log_info "操作已取消"
            exit 0
        fi
    else
        log_info "自动确认（非交互模式）"
    fi

    echo ""

    # 配置authorized_keys
    setup_authorized_keys "$target_user"

    echo ""

    # 配置SSH服务器
    configure_sshd

    echo ""

    # 测试配置
    if ! test_sshd_config; then
        log_error "配置测试失败，请检查配置文件"
        exit 1
    fi

    # 重启SSH服务
    restart_sshd

    # 显示摘要
    show_summary
}

# 执行主函数
main "$@"
